#include "usart.h"
#include "n32g430_usart.h"
#include <log.h>
#include "ota.h"
#include "misc.h"
#include <string.h>
#include "SEGGER_RTT.h"
#include "SEGGER_RTT_Conf.h"


#define USARTy            USART2
#define USARTy_CLK        RCC_APB1_PERIPH_USART2
#define USARTy_GPIO       GPIOA
#define USARTy_GPIO_CLK   RCC_AHB_PERIPH_GPIOA
#define USARTy_RxPin      GPIO_PIN_3
#define USARTy_TxPin      GPIO_PIN_2
#define USARTy_Rx_GPIO_AF GPIO_AF5_USART2
#define USARTy_Tx_GPIO_AF GPIO_AF5_USART2
#define USARTy_APBxClkCmd RCC_APB1_Peripheral_Clock_Enable
#define USARTy_IRQn       USART2_IRQn
#define USARTy_IRQHandler USART2_IRQHandler

// 接收缓冲区结构体
USART_RX_TypeDef USARTy_RX_Struct = {0};

void USARTy_Init(uint32_t baudrate)
{
    GPIO_InitType GPIO_InitStructure;
    USART_InitType USART_InitStructure;
  
    RCC_ClocksType RCC_ClocksStatus;

    // 使能时钟
    RCC_AHB_Peripheral_Clock_Enable(USARTy_GPIO_CLK);
    USARTy_APBxClkCmd(USARTy_CLK);
    RCC_Clocks_Frequencies_Value_Get(&RCC_ClocksStatus);

    // 配置GPIO
    GPIO_InitStructure.Pin = GPIO_PIN_2;  // TX
    GPIO_InitStructure.GPIO_Mode = GPIO_MODE_AF_PP;
    GPIO_InitStructure.GPIO_Alternate = USARTy_Tx_GPIO_AF;
    GPIO_Peripheral_Initialize(USARTy_GPIO, &GPIO_InitStructure);
  
    GPIO_InitStructure.Pin = GPIO_PIN_3;  // RX
    GPIO_InitStructure.GPIO_Mode = GPIO_MODE_AF_PP;
    GPIO_InitStructure.GPIO_Alternate = USARTy_Rx_GPIO_AF;
    GPIO_Peripheral_Initialize(USARTy_GPIO, &GPIO_InitStructure);
  
    // 配置USART
    USART_InitStructure.BaudRate = baudrate;
    USART_InitStructure.WordLength = USART_WL_8B;
    USART_InitStructure.StopBits = USART_STPB_1;
    USART_InitStructure.Parity = USART_PE_NO;
    USART_InitStructure.HardwareFlowControl = USART_HFCTRL_NONE;
    USART_InitStructure.Mode = USART_MODE_RX | USART_MODE_TX;
    USART_Initializes(USARTy, &USART_InitStructure);
  
    // 配置中断
    NVIC_InitType NVIC_InitStructure;

    NVIC_Priority_Group_Set(NVIC_PER2_SUB2_PRIORITYGROUP);
    /* Enable the USARTy Interrupt */
    NVIC_InitStructure.NVIC_IRQChannel                   = USARTy_IRQn;
    NVIC_InitStructure.NVIC_IRQChannelSubPriority           = 0;
    NVIC_InitStructure.NVIC_IRQChannelCmd                = ENABLE;
    NVIC_Initializes(&NVIC_InitStructure);
  
    // 使能接收中断
    USART_Interrput_Enable(USARTy, USART_INT_RXDNE);
    
    // 使能USART
    USART_Enable(USARTy);
}

// 发送单个字节
void USART_SendByte(uint8_t data)
{
    USART_Data_Send(USARTy, data);
    while(USART_Flag_Status_Get(USARTy, USART_FLAG_TXDE) == RESET);
}


void USART_SendBuffer(uint8_t *buffer, uint16_t size)
{
    for(uint16_t i = 0; i < size; i++) {
        USART_SendByte(buffer[i]);
    }
}

uint16_t USART_ReceiveBuffer(uint8_t *buffer, uint16_t size)
{
    uint16_t count = 0;
    uint32_t timeout = 0;
    
    // 等待接收完成或超时
    while(!USARTy_RX_Struct.rx_complete && timeout < 0xFFFF) {
        timeout++;
        delay(1);
    }
    
    if(USARTy_RX_Struct.rx_complete) {
        // 复制数据
        count = (USARTy_RX_Struct.rx_count > size) ? size : USARTy_RX_Struct.rx_count;
        memcpy(buffer, USARTy_RX_Struct.rx_buffer, count);
        
        // 清除接收状态
        USARTy_RX_Struct.rx_count = 0;
        USARTy_RX_Struct.rx_complete = 0;
    }
    
    return count;
}

void delay(uint32_t ms)
{
    uint32_t i;
    for(i = 0; i < ms * 10000; i++) {
        __NOP();
    }
}

// 升级检查函数
void USART_CheckUpgrade(void)
{
    if(USARTy_RX_Struct.rx_complete) {
        OTA_Packet *packet = (OTA_Packet*)USARTy_RX_Struct.rx_buffer;
        // 检查同步字和命令
        if(packet->sync == OTA_SYNC_WORD && packet->cmd == OTA_CMD_START) {
            LOGI("Received upgrade command\r\n");
            // 设置升级标志
            OTA_SetUpgradeFlag();
            LOGI("System will restart to upgrade\r\n");
            // 发送响应
            USART_SendByte(OTA_RESP_OK);
            delay(100);  // 等待串口消息发送完成
            NVIC_SystemReset();  // 重启进入bootloader
        }
        // 清除接收标志
        USARTy_RX_Struct.rx_count = 0;
        USARTy_RX_Struct.rx_complete = 0;
    }
}

// 计算校验和函数
static uint16_t Calculate_Packet_Checksum(uint8_t *data, uint16_t size)
{
    uint16_t sum = 0;
    for(uint16_t i = 0; i < size; i++) {
        sum += data[i];
    }
    return sum;
}


// 串口1中断服务函数
void USARTy_IRQHandler(void)
{
    if(USART_Flag_Status_Get(USARTy, USART_FLAG_RXDNE) != RESET) {
        uint8_t data = USART_Data_Receive(USARTy);
        
        switch(USARTy_RX_Struct.rx_state) {
            case RECV_SYNC1:
                if(data == (uint8_t)(OTA_SYNC_WORD & 0xFF)) {
                    USARTy_RX_Struct.rx_count = 0;
                    USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.rx_count++] = data;
                    USARTy_RX_Struct.rx_state = RECV_SYNC2;
                    LOGI("Sync1 received: 0x%02X\r\n", data);
                }
                break;
                
            case RECV_SYNC2:
                if(data == (uint8_t)(OTA_SYNC_WORD >> 8)) {
                    USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.rx_count++] = data;
                    USARTy_RX_Struct.rx_state = RECV_LENGTH;
                    LOGI("Sync2 received: 0x%02X\r\n", data);
                } else {
                    USARTy_RX_Struct.rx_state = RECV_SYNC1;
                    LOGI("Invalid Sync2: 0x%02X\r\n", data);
                }
                break;
                
            case RECV_LENGTH:
                USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.rx_count++] = data;
                if(USARTy_RX_Struct.rx_count == 4) {  // 接收完长度字段
                    uint16_t length = USARTy_RX_Struct.rx_buffer[2] | 
                                   (USARTy_RX_Struct.rx_buffer[3] << 8);
                    // 总长度 = sync(2) + length(2) + data(length) + checksum(2)
                    USARTy_RX_Struct.packet_length = length + 6;
                    USARTy_RX_Struct.data_length = length;
                    
                    // 检查长度是否合法
                    if(USARTy_RX_Struct.packet_length > USART_RX_BUFFER_SIZE) {
                        LOGI("Error: Packet too large! length=%d\r\n", length);
                        USARTy_RX_Struct.rx_state = RECV_SYNC1;
                        break;
                    }
                    
                    LOGI("Length field received: %d bytes\r\n", length);
                    USARTy_RX_Struct.rx_state = RECV_DATA;
                }
                break;
                
            case RECV_DATA:
                USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.rx_count++] = data;
                
                // 检查是否接收完数据部分
                if(USARTy_RX_Struct.rx_count == USARTy_RX_Struct.packet_length - 2) {
                    LOGI("Data received, entering checksum state\r\n");
                    USARTy_RX_Struct.rx_state = RECV_CHECKSUM;
                }
                break;
                
            case RECV_CHECKSUM:
                USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.rx_count++] = data;
                
                // 检查是否接收完整个数据包
                if(USARTy_RX_Struct.rx_count == USARTy_RX_Struct.packet_length) {
                    // 计算校验和
                    uint16_t calc_checksum = Calculate_Packet_Checksum(
                        USARTy_RX_Struct.rx_buffer,
                        USARTy_RX_Struct.packet_length - 2  // 不包括校验和
                    );
                    
                    // 获取接收到的校验和
                    uint16_t recv_checksum = 
                        (USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.packet_length - 1] << 8) |
                         USARTy_RX_Struct.rx_buffer[USARTy_RX_Struct.packet_length - 2];
                    
                    LOGI("Checksum: calculated=0x%04X, received=0x%04X\r\n",
                           calc_checksum, recv_checksum);
                    
                    if(calc_checksum == recv_checksum) {
                        LOGI("Packet complete and valid\r\n");
                        USARTy_RX_Struct.rx_complete = 1;
                        // 校验成功不需要立即发送响应，由OTA_ProcessCommand处理
                    } else {
                        LOGI("Checksum error! Sending error response\r\n");
                        // 发送错误响应，通知上位机重发
                        USART_SendByte(OTA_RESP_ERROR);
                        // 重置接收状态，准备接收重发的数据包
                        USARTy_RX_Struct.rx_count = 0;
                        USARTy_RX_Struct.rx_complete = 0;
                    }
                    
                    // 无论校验是否成功，都切换到等待同步字状态
                    USARTy_RX_Struct.rx_state = RECV_SYNC1;
                    
                    // 打印调试信息
                    OTA_Packet *packet = (OTA_Packet*)USARTy_RX_Struct.rx_buffer;
                    LOGI("Packet info: cmd=0x%02X, length=%d\r\n",
                           packet->cmd, packet->length);
                    
                    // 如果是数据包，打印地址信息
                    if(packet->cmd == OTA_CMD_DATA && packet->length >= sizeof(uint32_t)) {
                        OTA_DataPacket *data_packet = (OTA_DataPacket*)packet->data;
                        LOGI("Data packet address: 0x%08X\r\n", data_packet->addr);
                    }
                }
                break;
        }
        
        // 检查缓冲区溢出
        if(USARTy_RX_Struct.rx_count >= USART_RX_BUFFER_SIZE) {
            LOGI("Buffer overflow!\r\n");
            USARTy_RX_Struct.rx_state = RECV_SYNC1;
            USARTy_RX_Struct.rx_count = 0;
        }
        
        USART_Interrupt_Status_Clear(USARTy, USART_INT_RXDNE);
    }
}

